# finding all roots for multinomial logit 
# loosely based on predict method for mnlogit objects Contributed by: Florian Oswald, Univeristy College London
# Contributed by: Jose Guerra
roots.social.mnlogit <- function(object, newdata=NULL, probability=FALSE, wdata=NULL, p0=NULL, ...)
{
  size <- object$model.size
  # get choice set for colnames
  choiceVar <- object$call$choiceVar
  choiceSet <- unique(index(object)$alt)#unique(object$data[[choiceVar]])
  if (is.null(wdata)) {
    stop("wdata containing weights must be provided")
  }
  if (is.null(p0)) {
    stop("initial probabilities must be provided for FXP iteration")
    #p0 <- matrix(cbind(1-rowSums(object$probabilities),object$probabilities), ncol = 1, byrow=FALSE)
  }
  if (is.null(newdata)) {
    # if no new data, use probabilities computed during training model
    if (probability)
      return(object$probabilities)
    else {
      return(apply(object$probabilities, 1, function(x)
        object$choices[which(x == max(x, na.rm = TRUE))]))
    }
  } else {
    # check that all columns from data are present
    # this is important when you build Y below.
    newn <- names(newdata)
    oldn <- names(object$data)
    if (!all(oldn %in% newn))
      stop("newdata must have same columns as training data. ")
    # Define initial probabilities
    p0 <- matrix(p0, ncol = 1, byrow=FALSE)
    if((ncol(p0)*nrow(p0)) %% size$K) 
      stop("Mismatch in p0, (nrow,ncol) must be either (N*L,1) or (N*L)")
    probP <- p0
    # make sure newdata is ordered by choice
    newdata <- newdata[order(newdata[[choiceVar]]), ]
    # different model size: N # newdata must have N*K rows
    if (nrow(newdata) %% size$K) #stands for modulus i.e. whether is divisible (0) or not (1)
      stop("Mismatch between number of rows in newdata and number of choices.")
    #check length of weighting matrix coincides with length of data
    if (nrow(newdata)/size$K != length(wdata$neighbours))
      stop("Mismatch between number of rows in newdata and number of rows in weighting matrix")
  }
 
  data <- newdata
  weigh <- wdata
  size$N <- nrow(data)/size$K # number of individuals
  cheatSocial <- unique(data$IDSocial)
  data[sapply(c(1:size$K),function(x) (x-1)*size$N+1),]$IDSocial <- rep(cheatSocial-1,size$K)
  cheatCivil <- unique(data$IDCivil)
  data[sapply(c(1:size$K),function(x) (x-1)*size$N+1),]$IDCivil <- rep(cheatCivil-1,size$K)
  # Initialize utility matrix: dim(U) = N x K-1
  probMat <- matrix(rep(0, size$N * (size$K-1)), nrow = size$N, ncol = size$K-1)
  formDesignMat <- function(varVec = NULL, includeIntercept = TRUE)
  {
    if (is.null(varVec) && !includeIntercept) return(NULL)
    fm <- paste(attr(formula, "response"), "~")
    if (!is.null(varVec))
      fm <- paste(fm, paste(varVec, collapse = "+"))
    if (!includeIntercept) fm <- paste(fm, "-1 ")
    else fm <- paste(fm, "+1 ")
    modMat <- model.matrix(as.formula(fm), data)
  }
  # Grab the parsed formula from the fitted mnlogit object
  #formula <- object$formula
  formula <- pparseFormula(object$formula)

  X <- formDesignMat(varVec = attr(formula, "indSpVar"),
                     includeIntercept = attr(formula, "Intercept"))
  X <- if (!is.null(X)) X[1:size$N, , drop=FALSE] # Matrix of ind sp vars
  X[1,grep("factor+",colnames(X))] <- 1 #to take back cheating
  Y <- formDesignMat(varVec = attr(formula, "csvChCoeff"),
                     includeIntercept = FALSE)
  Z <- formDesignMat(varVec = attr(formula, "csvGenCoeff"),
                     includeIntercept = FALSE)
  # Do the subtraction: Z_ik - Zi0 (for Generic coefficients data)
  ### NOTE: Base choice (with respect to normalization) is fixed here
  ### Base choice is the FIRST alternative
  if(!is.null(Z)) {
    for (ch_k in 2:size$K) {
      Z[((ch_k - 1)*size$N + 1):(ch_k*size$N), ] <-
        Z[((ch_k-1)*size$N+1):(ch_k*size$N), , drop=FALSE]
      - Z[1:size$N, , drop=FALSE]
    }
  }
  # Drop rows for base alternative
  Z <- Z[(size$N + 1):(size$K*size$N), , drop=FALSE]
  # Grab trained model coeffs from fitted mnlogit object
  coeffVec <- object$coeff
  matches <- unique( grep(paste(paste(c("\\(Intercept\\)",colnames(X),paste("factor\\(IDSocial\\)",cheatSocial,sep="")),":+",sep=""),collapse="|"), names(coeffVec), value=TRUE))  
  matchesY <- unique(grep(paste(paste(colnames(Y),":+",sep=""),collapse="|"), names(coeffVec), value=TRUE))  
  coeffVecY <- coeffVec[matchesY]
  #coeffVecY[c(1:length(coeffVecY))] <- runif(length(coeffVecY),3,5)
  coeffVec <- coeffVec[matches] #subsetting coefficients per parish
  nmiss <-0  
  if (length(grep("factor\\(IDSocial\\)+",names(coeffVec)))==0)  {
    coeffVec <- c(coeffVec,rep(0,(size$K-1)))
    names(coeffVec)[(length(coeffVec)-(size$K-1)+1):length(coeffVec)] <- paste("factor(IDSocial)",cheatSocial,":",c(1:(size$K-1)),sep="")
    nmiss <- nmiss+1
    coeffVec <- coeffVec[c(names(coeffVec)[1:(length(coeffVec)-nmiss*(size$K-1))],
                           sort(names(coeffVec)[(length(coeffVec)-nmiss*(size$K-1)+1):length(coeffVec)]))]
  }
  neword <- c(eval(parse(text=paste("grep(","'+:",1,"'",",","names(coeffVec))",sep=""))))
  size$p <- length(neword) #-size$f #number of ind specific coefficients
  #for(i in c(1:(size$K-1))) {
  for(i in c(2:(size$K-1))) {
    neword <- c(neword,eval(parse(text=paste("grep(","'+:",i,"'",",","names(coeffVec))",sep=""))))
  }
  coeffVec <- c(coeffVec[neword],coeffVecY) #reordering the coefficients so matrix multiplication is comformable
  coeffVec[is.na(coeffVec)]<-0 #for those that were given an NA
  
  # First compute the utility matrix (stored in probMat)
  if (size$p) {
    probMat <- probMat + X %*% matrix(coeffVec[1:((size$K-1) *size$p)],
                                      nrow = size$p, ncol = (size$K-1), byrow=FALSE)
  }
  
  #update data by weighting 
  weightprobMat<- function(ch_k)
  {
    init <- (ch_k - 1)*size$N + 1
    fin <- ch_k * size$N
    probPprov <- probP[init:fin, drop=FALSE]
    #attr(probMatprov,"class") <- "data.frame"
    lag.listw(wdata,probPprov)
  }
  for(ch_k in c(1:size$K)) {
    init <- (ch_k - 1)*size$N + 1
    fin <- ch_k * size$N
    Y[init:fin, drop=FALSE] <- weightprobMat(ch_k)
  }
  #newdatacopy <<- newdatacopy  
  
  if (size$f) {
    findYutil<- function(ch_k)
    {
      offset <- (size$K - 1)*size$p
      init <- (ch_k - 1)*size$N + 1
      fin <- ch_k * size$N
      Y[init:fin, , drop=FALSE] %*%
        coeffVec[((ch_k-1)*size$f + 1 + offset):(ch_k*size$f+offset)]
    }
    vec <- as.vector(sapply(c(1:size$K), findYutil))
    
    # normalize w.r.t. to k0 here - see vignette on utility normalization
    vec <- vec - vec[1:size$N]
    probMat <- probMat + matrix(vec[(size$N+1):(size$N*size$K)],
                                nrow = size$N, ncol = (size$K-1), byrow=FALSE)
  }
  
  if (size$d) {
    probMat <- probMat +
      matrix(Z %*% coeffVec[(size$nparams - size$d + 1):size$nparams],
             nrow = size$N, ncol=(size$K-1), byrow=FALSE)
  }
 
  # Convert utility to probabilities - use logit formula
  probMat <- exp(probMat) # exp(utility)
  baseProbVec <- 1/(1 + rowSums(probMat)) # P_i0
  probMat <- probMat * matrix(rep(baseProbVec, size$K-1),
                              nrow = size$N, ncol = size$K-1) # P_ik
  probMat <- cbind(baseProbVec,probMat)
  ProbMat <<- probMat

  if (nrow(probMat) == 1)
    probMat <- as.matrix(probMat)
  colnames(probMat) <- choiceSet
  if (probability)
    #return(probMat)
    return(c(matrix(probMat, ncol = 1, byrow=FALSE)))
  else {
    choice <- apply(probMat, 1, function(x)
      object$choices[which(x == max(x, na.rm = TRUE))])
    return(choice)
  } 
}

